#pragma once

#include <boost/unordered_map.hpp>
#include "MDPSolver.h"
#include <list>
using namespace std; 


/*
	This class discretizes the state space and solves the resulting finite MDP using 
	value iteration. To build the transition function, it samples several points from each 
	cell and follows the transition function to see which cell it lands in.
*/



namespace discreteSolver
{

class DiscreteMDPSolver; 
class DiscretizedState; 

struct TransProb
{
public: 
	DiscretizedState* nextState; 
//	long nextIndex;			removing for performance improvement
	double prob; 
	double reward; 
};

/*
	This class holds information about one discretized state. 
*/
class DiscretizedState
{
public:
	DiscretizedState(long si, DiscreteMDPSolver* p); 
	~DiscretizedState(); 
	void update();				//do one bellman backup
	double getValue();  
	double getQValue(int action){return qvals[action];}
	void makeRandomState(Observation o);	//generate a random state in this cell (o should already be allocated)
	Action getBestAction();					

	void print(FILE* f);					

	long  m_index;			//my index in the entire state space
	
	double* mins;				//the boundary of this state
	double* maxs; 

	list<TransProb>* tfunc;		//this is the transition function (an array of action elements, each having the prob of next state)
	double* qvals;
	double value;				//we should make sure value is always synced with qvals
	DiscreteMDPSolver* parent; 
};

// This class builds the transition function for a batch of states (indicated by start-end)
class FiniteMDPTFBuilder
{
public:
	FiniteMDPTFBuilder(DiscreteMDPSolver* p, long minInd, long  maxInd){parent = p; start = minInd; end = maxInd; }
	void operator()(){buildTF();}
	void buildTF(); 
	
public:
	DiscreteMDPSolver* parent; 
	long start;					//the start of the index we should fill the TF for 
	long  end;					//the end index we should fill (this is exclusive)
};


/*
	This class runs in background and does value iteration by constantly picking a random state and bellman updating it
*/
class FiniteMDPVI
{
public:
	FiniteMDPVI(DiscreteMDPSolver* p){parent = p; }
	void operator()(){doVI();}
	void doVI(); 

public:
	DiscreteMDPSolver* parent; 
}; 



class DiscreteMDPSolver : public SlaveableMDPSolver
{
public:
	DiscreteMDPSolver(MREAgent* a, TFGenerator* tf, RFGeneralizer* rf);
	virtual ~DiscreteMDPSolver();

	//general required methods for MDPSolver:
	virtual void operator ()(); 
	virtual Action getBestAction(Observation state); 
	virtual void	solveModel(Observation currentState); 
	virtual double getStateValue(const Observation o); 

	//specific to this class
	long observation2Index(const Observation o); 
	void constructStates();					//creates the discretization
	void destroyStates();					
	void clearStateContents();				//clear the transition function from states (usefull for replanning where we don't want to recreate the states themselves)
	void constructTransitionFunction(); 
	void doVI(); 

	friend class FiniteMDPVI;				//so we can do multi-threaded VI
	friend class FinitMDPTFBuilder;			//so we can make the transition function in multi-thread

public:
	DiscretizedState** states;			//this is an array of pointers
	long stateSize; 
	bool parentTFRunning;					//when dealing with slaves in thread, this is how we tell the children to stop computing
	bool parentVIRunning; 
}; 

}//namespace


